Skip to content

feat : a general probabilistic forecasting interface#700

Open
Sir-Sloth-The-Lazy wants to merge 17 commits into
mllam:mainfrom
Sir-Sloth-The-Lazy:feat/probabilistic-forecasting-interface
Open

feat : a general probabilistic forecasting interface#700
Sir-Sloth-The-Lazy wants to merge 17 commits into
mllam:mainfrom
Sir-Sloth-The-Lazy:feat/probabilistic-forecasting-interface

Conversation

@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor

Describe your changes

This PR implements the abstract constructions agreed in the #685 discussion (RFC: a general probabilistic forecasting interface), deliberately without any concrete Graph-EFM instantiation, so the design can be reviewed on its own and this can server as the base for addition of any probabilistic model, not just graph_efm.

Move ownership of the training objective from ForecasterModule onto the Forecaster. A new abstract Forecaster.compute_training_loss returns a finished (loss, loss_components) pair, so each forecaster owns its complete training objective (including assembling it from any internal terms).ForecasterModule.training_step now only injects the configured scoring rule (--loss), the interior mask and per_var_std, then logs the returned loss and components (component names prefixed with the phase). Per the discussion,
this move is applied to the deterministic setup as well, so the same concept sits in the same place in both stacks: the deterministic ARForecaster training loss is unchanged in value (covered by an equality test), it is just computed by the forecaster itself. Note for reviewers: since the method is abstract on Forecaster, any out-of-tree Forecaster subclass now has to implement it; all in-repo forecasters go through ARForecaster.

Add the probabilistic side, built on top of the deterministic one (no sample dimension leaks into deterministic components):

  • ProbabilisticForecaster (abstract): declares sample_ensemble, producing members stacked along a new dimension after batch, (B, S, pred_steps, num_grid_nodes, d_state). This encodes the module's only assumption "the forecaster can create ensemble forecasts of the correct shape" as a type contract, and leaves room for non-AR implementations (diffusion/flow) later.

  • ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): unrolls independent trajectories through a step predictor that samples its output (sequential loop as the safe first default), and by default trains on the configured scoring rule applied to the ensemble mean. Model-specific objectives (e.g. Graph-EFM's ELBO + CRPS) will live in subclasses that override compute_training_loss.

  • ProbabilisticForecasterModule: training inherited unchanged; validation samples an eval_ensemble_size ensemble and logs the RMSE of the ensemble mean (the example ensemble metric suggested in RFC: a general probabilistic forecasting interface #685, standing in until the ensemble-metrics PR), reusing the val_mean_loss / val_loss_unroll{i} log keys so existing checkpoint callbacks work unchanged. test_step raises NotImplementedError rather than silently evaluating a single member deterministically; ensemble test evaluation and plotting are a follow-up.

Dependencies: none. In particular this PR does not depend on the planned ensemble metrics (crps_ens, spread_squared). Follow-ups: ensemble metrics, the Graph-EFM forecaster (ELBO + CRPS, its own kl_beta/crps_weight config), ensemble test evaluation + plotting, and train_model.py wiring for
probabilistic models.

Issue Link

Type of change

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README to cover introduced code changes (no README changes needed — no user-facing CLI/config changes in this PR)
  • I have added tests that prove my fix is effective or that my feature works
  • I have given the PR a name that clearly describes the change, written in imperative form (context).
  • I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section
    reflecting type of change (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

Add abstract Forecaster.compute_training_loss returning a finished
(loss, loss_components) pair, so each forecaster owns its complete
training objective. ForecasterModule.training_step now only injects the
configured scoring rule, interior mask and per_var_std, and logs the
result. The deterministic ARForecaster loss is unchanged in value.

Add the abstract ProbabilisticForecaster (sample_ensemble capability),
ProbabilisticARForecaster (sequential sampled rollouts, trains on the
configured score of the ensemble mean) and a minimal
ProbabilisticForecasterModule whose validation samples an ensemble and
logs the RMSE of the ensemble mean. Interface design from mllam#685.
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

@joeloskarsson @observingClouds , this is the implementation of the dicussion on issue #685 hope this is what you wanted ! 👀.

@Sir-Sloth-The-Lazy
Sir-Sloth-The-Lazy marked this pull request as draft July 5, 2026 04:20

@observingClouds observingClouds left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @Sir-Sloth-The-Lazy thanks for drafting this. I had a look at the discussion around the proposal in #685 and this PR. It looks well aligned. So far, I only have a few minor comments.

trajectory. This class adds ensemble forecasting on top: unrolling
several trajectories and stacking them along an ensemble dimension.
The default training objective scores the ensemble mean with the
injected scoring rule; forecasters with model-specific objectives

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"injected scoring rule". Can this be written more explicitly? So that it is clear to users where the scoring rule is set.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

solved in commit 3f5402d

states at each predicted step, used both as the prediction
targets and to overwrite boundary nodes during the rollouts.
Dims: same as one ensemble member.
score_fn : Callable

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fn always reminds me as an abbreviation for filename. Maybe use score_func or score_metric?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Or even just score.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Solved in commit 3f5402d

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

renamed to score_metric as per the suggestion.

@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

@observingClouds , Hope the latest commit solves the concerns 😁

@joeloskarsson joeloskarsson left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks quite nice I think, and is still pretty easy to follow. I had some outstanding thoughts regarding the connection to the deterministic version + some small things.

Comment thread neural_lam/models/module.py Outdated
init_states,
forcing_features,
target_states,
score_metric=self.loss,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the loss calculation is now entirely handled by the Forecaster, I think it could make more sense that self.loss is instantiated and sits on the Forecaster, rather than being passed to it? What do you think, is there any problem with that?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question, I don't think there's a hard problem with it, but it's a bit bigger than a one-liner so wanted to lay out the scope before doing it.

Right now self.loss isn't only feeding compute_training_loss ForecasterModule also calls it directly in validation_step/test_step (and for the spatial loss map) to report val/test loss, independent of whatever the forecaster's training objective ends up being. So "loss calculation entirely handled by the Forecaster" isn't quite true yet for the deterministic path the Module still owns a second, separate use of the same metric for reporting.

It also isn't uniform across the two modules we have: ProbabilisticForecasterModule.validation_step doesn't touch self.loss at all, it scores the ensemble mean with raw metrics.mse, and test_step isn't implemented there yet.

If we move it, the plan would be:

loss: str becomes a ctor arg on Forecaster (or ARForecaster), instantiated there as self.loss = metrics.get_metric(loss)
compute_training_loss drops the score_metric param entirely and reads self.loss internally
ForecasterModule.validation_step/test_step read self.forecaster.loss instead of self.loss
loss gets threaded through the forecaster constructors in train_model.py (currently only passed to ForecasterModule) and the ~15 test call sites that build ARForecaster/ProbabilisticARForecaster directly
None of that is a blocker, just confirming that's the scope before I make the change and whether ForecasterModule should keep a loss kwarg at all (for backward-compat / overriding the reporting metric independently of the training metric) or drop it entirely in favor of always reading it off the forecaster.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not a fan of having the ForecasterModule get the loss as self.forecaster.loss and then using that. A correct division of responsibilities would mean that the ForecasterModule does not know how to compute the loss for a batch, and that should be delegated to the Forecaster. So instead of getting the loss and using it, I think any loss computation that is needed in the validation/test step should be delegated to the Forecaster. If that means that we need additional functionality in the Forecaster (e.g. other types of loss reduction, not all dims) than we should implement that there.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

f0e01b1 committed the fix

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

f0e01b1 committed the fix

Comment thread neural_lam/models/module.py Outdated
target_states,
score_metric=self.loss,
interior_mask_bool=self.interior_mask_bool,
per_var_std=self.per_var_std,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly for per_var_std, the main use case of this seems to be computing the loss, so would it also belong in the Forecaster then, that handles the loss?

@Sir-Sloth-The-Lazy Sir-Sloth-The-Lazy Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • On score_metric/self.loss: Moved. ARForecaster (and ProbabilisticARForecaster through it) now takes a loss: str = "wmse" constructor argument and builds self.loss = metrics.get_metric(loss) itself. compute_training_loss no longer takes a score_metric param, it reads self.loss internally. ForecasterModule.training_step now just injects the interior mask; validation_step/test_step read self.forecaster.loss for their reporting metric instead of owning a separate copy.

  • On per_var_std: Moved too, same reasoning. ARForecaster also takes an optional config: NeuralLAMConfig = None constructor argument and computes self.per_var_std from it (same formula as before diff_std / sqrt(feature_weights)), only when not self.predicts_std. Made config optional rather than required: since per_var_std is only needed for scoring, a forecaster built purely for inference (or the several existing tests that only exercise forward()) doesn't need to supply a NeuralLAMConfig at all it stays None and is simply never used, same as today's predicts_std=True case.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

made the above change in 987fecd

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This sounds good, but what happens if we use StepPredictor that does not output a pred_std, but also no NeuralLAMConfig is given as an argument ( = None)? That seems to potentially cause issues.
I am not certain if we need to allow the config=None case, there should always be such a config present so we could require it.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see you noticed this as well, and fixed in 511a6d5. However, I don't think this check sits in the right spot. It checks the Forecaster configuration, so it should be in the Forecaster rather than the ForecastModule probably.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pushed this as 8b3ab33

Comment thread neural_lam/models/module.py Outdated
Comment thread neural_lam/models/forecasters/probabilistic.py Outdated
Comment thread neural_lam/models/forecasters/probabilistic.py
Comment thread neural_lam/models/forecasters/probabilistic.py Outdated
Comment thread neural_lam/models/probabilistic_module.py Outdated
Comment thread neural_lam/models/probabilistic_module.py Outdated
Comment thread neural_lam/models/probabilistic_module.py Outdated
Comment thread CHANGELOG.md Outdated
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

Thank you for taking out the time to review this @joeloskarsson means a lot. I will start working on the review now 😃

Sir-Sloth-The-Lazy and others added 7 commits July 8, 2026 09:10
score_metric/per_var_std were injected into compute_training_loss by
ForecasterModule and also used directly for val/test loss reporting,
duplicating config the forecaster already needs for its own objective.
ARForecaster/ProbabilisticARForecaster now own self.loss and
self.per_var_std (computed from an optional config ctor arg), and
ForecasterModule reads them off self.forecaster instead. Also trims
the CHANGELOG entry for mllam#685 down to one sentence per review feedback.
A Forecaster built without config now silently has per_var_std=None
when its predictor doesn't output its own std. Previously per_var_std
was always computed by ForecasterModule itself, so this gap didn't
exist; now that construction is split across two calls, catch it at
ForecasterModule init instead of crashing at the first val/test step.
Each member's predicted std is its own, not a spread computed across
the ensemble, so ensemble_std was a misleading name. Document on
ProbabilisticForecaster that a per-member std makes the predictive
distribution a mixture of Gaussians, and note in
ProbabilisticARForecaster.compute_training_loss that averaging the
per-member stds is a simplification of the true mixture variance
(which also includes the spread between member means).
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Scoring the ensemble mean with a pointwise metric only rewards the
mean being right, giving the model no incentive to keep a calibrated
spread, and risks training it to collapse to a point estimate. Redeclare
compute_training_loss as abstract on ProbabilisticARForecaster instead
of providing that as a default (it would otherwise silently fall back
to ARForecaster's single-rollout objective via MRO, not even the
ensemble mean). Concrete subclasses must define their own objective.

Tests that only need an instantiable forecaster now use a local
ConcreteProbabilisticARForecaster example (ensemble-mean scoring,
moved out of the library code); a new test locks in that the base
class itself cannot be instantiated.
Drop ProbabilisticARForecaster's ensemble_size constructor arg and the
implicit num_members=None -> self.ensemble_size fallback in
sample_ensemble; num_members is now always required. Baking a default
member count into the forecaster's state was unnecessary now that
compute_training_loss is abstract too (nothing in the shared base
class path used it) and just adds an implicit default callers could
silently rely on instead of deciding explicitly. The num_members < 1
validation moves from __init__ to sample_ensemble accordingly.
ProbabilisticForecasterModule.eval_ensemble_size follows suit: it no
longer defaults to None with a forecaster fallback, it's required.

Test-only ConcreteProbabilisticARForecaster (used wherever a concrete
probabilistic forecaster is needed for testing) gains its own
train_num_members for the training objective, since deciding how many
members to sample during training is now the concrete subclass's call.
Mirrors validation_step: samples eval_ensemble_size members and scores
the ensemble mean, same as validation. Factored the shared sampling +
scoring + logging into _ensemble_step(batch, phase) rather than
duplicating the block, since validation_step and test_step differ only
in their log-key prefix and which metrics dict collects the result.

Overrides on_test_epoch_end (rather than inheriting
ForecasterModule's) since this module's test_step doesn't populate
spatial_loss_maps or plot examples - the inherited version would crash
on torch.cat of an empty list.
Comment thread tests/test_probabilistic_forecaster.py Outdated
Comment on lines +64 to +65
# per_var_std is normally computed from config; override directly since
# this test only cares about the loss computation, not standardization.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not much of a fan of these inline comments and think the code should be clear enough to explain itself.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair point, removed the assignment is clear on its own. Committed as 45ffdeb. Is there something more you want out of this ? I would really like to see, if this is enough or I can do more to make the code more readable.

from .module import ForecasterModule


class ProbabilisticForecasterModule(ForecasterModule):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of using inheritance here I would make an abstract base class, e.g. BaseForecasterModule, which implements expected methods as abstractmethods, and then renaming the ForecasterModule as DeterministicForecasterModule. And then ProbabilisticForecasterModule and DeterministicForecasterModule become an implementation of this base class.

Doing it this way will make the separation of concerns clearer, and makes it clear what any kind of "forecaster" must behave like (take as input, give as output)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would put these in .forecasters.base_module, .forecasters.deterministic_module and .forecasters.probablistic_module

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

any shared functionality between the deterministic and probabilistic forecaster modules would then sit in the base forecaster module

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am wondering whether "Forecaster" is actually a good term as this limits the usage to forecasts, but we might also have negative timestepping at some point or do downscaling. So is "Predictor" maybe a better/more general term here?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Restructured the module hierarchy per @leifdenby's suggestion:

  • Introduced BaseForecasterModule(pl.LightningModule, ABC) in neural_lam/models/forecasters/base_module.py, holding everything shared between evaluation modes: training_step, common_step, on_after_batch_transfer (batch standardization), _warn_skipped_val_steps, on_validation_epoch_end, plot_examples, aggregate_and_plot_metrics, on_load_checkpoint. validation_step, test_step, and on_test_epoch_end are left @abstractmethod since they differ meaningfully between deterministic and ensemble evaluation.

  • Renamed ForecasterModuleDeterministicForecasterModule and moved it to neural_lam/models/forecasters/deterministic_module.py.

  • Moved ProbabilisticForecasterModule to neural_lam/models/forecasters/probabilistic_module.py. It now implements BaseForecasterModule directly as a sibling of the deterministic module, rather than subclassing it. This also resolves the tension from the earlier super().validation_step() discussion since both modules independently implement the same abstract contract, there's no more inheritance-reuse temptation (and no risk of wasting a forward pass or reintroducing the loss/diagnostic-naming conflation that came up earlier in review).

All 256 tests and pre-commit hooks (black/isort/flake8/interrogate/mypy) pass. Touched 16 files total (mostly mechanical import/rename updates in tests and docstrings).

Commit: d020691.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, worth thinking through. Two observations that make me lean toward keeping "Forecaster" for now, though I'm open to it:

  1. There's already a StepPredictor in this same package the single-step model that ARForecaster wraps and unrolls into a multi-step trajectory. Right now "Forecaster wraps StepPredictor" reads unambiguously as two different levels of abstraction. If Forecaster becomes Predictor, we'd have Predictor (multi-step) sitting next to StepPredictor (single-step), the containment relationship gets muddier, since the names no longer signal which one is "bigger."

  2. Backward timestepping and downscaling are worth keeping the door open for, but I think even those still fit "produce a forecast/estimate over a target domain from some initial condition" reasonably well. "Predictor" being more generic also makes it vaguer about what is being predicted, which the current interface (states at future grid times) is actually fairly specific about.

So unless we're also willing to rename StepPredictor at the same time to keep the two levels distinct (e.g. Forecaster/StepModel), I'd rather not do a naming-only rename across the whole package without a concrete near-term use case (downscaling, hindcasting) driving it. Happy to revisit this once one of those lands.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a very good suggestion, fully agree. Also good to keep the classes in separate files.

This also resolves the tension from the earlier super().validation_step() discussion since both modules independently implement the same abstract contract

Indeed, that's very nice!

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As discussed a bit on slack, it seems best to wait with any renaming of these classes and stick to Forecaster now.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also think the files containing the ForecasterModules should be in their own modules directory, and not mixed with the forecasters.

@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

Thank you for your review ! I am on it 🤩

Rename the ensemble-mean diagnostic keys from *_loss_unroll/*_mean_loss
to *_ens_rmse_unroll/*_mean_ens_rmse so they aren't conflated with the
training loss, per review feedback.
Remove explanatory comments around the per_var_std overrides; the
assignments are clear on their own.
…crete deterministic/probabilistic modules

Introduce BaseForecasterModule (abstract) under models/forecasters/ holding
shared plumbing (training_step, common_step, batch standardization,
checkpoint compatibility, plotting/aggregation helpers), with
validation_step, test_step and on_test_epoch_end left abstract since they
differ meaningfully between evaluation modes. Rename ForecasterModule to
DeterministicForecasterModule and move it, alongside
ProbabilisticForecasterModule, into forecasters/ as siblings implementing
the shared contract, rather than one subclassing the other.
@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor Author

@leifdenby @joeloskarsson @observingClouds, I hope these commits solve most of the issue, where i need a little more direction, I have left the comments asking for more help ! Thank you hope to here from you all soon. 🤩

Sir-Sloth-The-Lazy added a commit to Sir-Sloth-The-Lazy/neural-lam that referenced this pull request Jul 14, 2026
Removes estimate_likelihood/compute_step_loss and the per_var_std
buffer they existed to feed, plus the now-unused config constructor
arg (its only use was building per_var_std). Per the objective now
living on the Forecaster (mllam#700), the predictor stays a pure network
construct: encoder/prior/decoder plus the forward sampling path.
@Sir-Sloth-The-Lazy
Sir-Sloth-The-Lazy marked this pull request as ready for review July 15, 2026 08:01
@joeloskarsson

Copy link
Copy Markdown
Collaborator

Read through all of the fixes and discussion here now. Some small follow-up comments on my earlier points, but in general I think this is looking quite good. The refactor into a base module class was very good, but also moved some things around a bit. I would be happy to give this another full read through before approving once I am back from vacation. I think that should be timely if this goes in in v0.8.0.

Add JetBrains IDE project directory to the ignore list alongside
the existing .vim/.vscode entries.
…aster

DeterministicForecasterModule.validation_step/test_step read
self.forecaster.loss and self.forecaster.per_var_std directly, so the
Module still knew how to compute a loss from a prediction. Add an
abstract Forecaster.score (implemented on ARForecaster) that resolves
the pred_std fallback and applies a scoring rule internally; the Module
now only calls forecaster.score(...) and never touches loss/per_var_std
itself.
…ster

BaseForecasterModule.__init__ raised if forecaster.per_var_std was None
and the forecaster didn't predict its own std -- a check on the
Forecaster's configuration living in the wrong class. Move it into
ARForecaster via a new _resolve_pred_std helper, shared by score() and
compute_training_loss(): construction with config=None now always
succeeds (a valid state for forecasters only ever used for inference),
and the ValueError instead fires from the Forecaster itself, only once
scoring is attempted without any std to use.
…dules/

BaseForecasterModule, DeterministicForecasterModule and
ProbabilisticForecasterModule (Lightning wrappers around a Forecaster)
were mixed in with the Forecaster classes themselves under
neural_lam/models/forecasters/. Move them into their own
neural_lam/models/modules/ package (base.py, deterministic.py,
probabilistic.py) so the two concerns: what a forecaster is vs how
it's trained/evaluated by Lightning live in separate directories.
Pure move: internal imports and neural_lam/models/__init__.py updated
accordingly, no behavioural change.
Resolves the conflict on neural_lam/models/module.py: main's PR mllam#675
added --train_steps_to_log and deduplicated the val/test/train
prediction+loss+logging pattern into _compute_prediction_and_loss /
_log_step_loss, on top of the pre-mllam#700 module.py where the Module
still owned self.loss/self.per_var_std directly. This branch had since
deleted that file (split into modules/base.py, deterministic.py,
probabilistic.py, with loss ownership moved onto the Forecaster).

Ports mllam#675's dedup/logging helpers (_warn_skipped_steps, _log_step_loss,
the train_steps_to_log hparam) into modules/base.py, and rewires
DeterministicForecasterModule's validation_step/test_step to use them
via forecaster.score() instead of the removed self.loss/self.per_var_std.
_compute_prediction_and_loss is dropped rather than ported: it assumes a
single per-step-decomposable loss, which conflicts with
compute_training_loss's contract that a forecaster's training objective
may not decompose per step (e.g. an ELBO). training_step therefore still
only logs the aggregate train_loss; train_steps_to_log is accepted for
CLI/checkpoint compatibility but does not yet produce a
train_loss_unroll{i} breakdown -- documented inline as a deliberate
scope boundary, not an oversight.

Also fixes two call sites (deterministic.py, probabilistic.py) that
referenced the pre-rename _warn_skipped_val_steps, and absorbs the
unrelated utils.py -> utils/ package split from mllam#682 (auto-merged
cleanly, this branch never touched utils.py).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants